// server_hub.go
//
// 并发安全特性：
//   - 使用 hub 管理所有客户端连接（注册、注销、广播）
//   - safeConn 为每个连接提供写锁，防止多协程同时 Write
//   - context 控制关闭时的协作，WaitGroup 等待 goroutine 退出
package main

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"net"
	"sync"
	"time"
)

// safeConn 封装 net.Conn，提供写锁保障写操作串行化。
type safeConn struct {
	net.Conn
	mu sync.Mutex // 写锁，保证一次仅有一个 goroutine 写 socket
}

func (sc *safeConn) SafeWrite(b []byte) (int, error) {
	sc.mu.Lock()
	defer sc.mu.Unlock()
	return sc.Conn.Write(b)
}

// hub 负责注册/注销连接以及广播消息。
type hub struct {
	mu        sync.RWMutex
	conns     map[string]*safeConn
	broadcast chan string
}

// newHub 初始化连接管理器。
func newHub() *hub {
	return &hub{
		conns:     make(map[string]*safeConn),
		broadcast: make(chan string, 32),
	}
}

// register 将连接加入管理。
func (h *hub) register(id string, sc *safeConn) {
	h.mu.Lock()
	h.conns[id] = sc
	h.mu.Unlock()
	log.Printf("[hub] register %s (total=%d)", id, h.size())
}

// unregister 从管理器移除连接。
func (h *hub) unregister(id string) {
	h.mu.Lock()
	if sc, ok := h.conns[id]; ok {
		delete(h.conns, id)
		sc.Close()
	}
	h.mu.Unlock()
	log.Printf("[hub] unregister %s (total=%d)", id, h.size())
}

// size 返回当前连接数量。
func (h *hub) size() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.conns)
}

// runBroadcast 循环消费广播通道，并向所有连接安全写入。
func (h *hub) runBroadcast(ctx context.Context, wg *sync.WaitGroup) {
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			select {
			case <-ctx.Done():
				log.Println("[hub] broadcast loop exit")
				return
			case msg := <-h.broadcast:
				h.mu.RLock()
				for id, sc := range h.conns {
					if _, err := sc.SafeWrite([]byte(msg)); err != nil {
						log.Printf("[hub] write to %s error: %v", id, err)
					}
				}
				h.mu.RUnlock()
			}
		}
	}()
}

// submit 广播一条消息。
func (h *hub) submit(msg string) {
	select {
	case h.broadcast <- msg:
	default:
		log.Println("[hub] broadcast queue full, drop message")
	}
}

const (
	serverAddr = "127.0.0.1:8000"
)

func main() {
	listener, err := net.Listen("tcp", serverAddr)
	if err != nil {
		log.Fatalf("listen error: %v", err)
	}
	defer listener.Close()
	log.Printf("[server] listening on %s", serverAddr)

	h := newHub()
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	h.runBroadcast(ctx, &wg)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("[server] accept error: %v", err)
			continue
		}
		sc := &safeConn{Conn: conn}
		id := conn.RemoteAddr().String()
		h.register(id, sc)

		wg.Add(1)
		go func(id string, sc *safeConn) {
			defer wg.Done()
			defer h.unregister(id)

			reader := bufio.NewScanner(sc)
			for {
				if err := sc.SetReadDeadline(time.Now().Add(60 * time.Second)); err != nil {
					log.Printf("[server] set read deadline error: %v", err)
					return
				}
				if !reader.Scan() {
					if err := reader.Err(); err != nil {
						log.Printf("[server] read error %s: %v", id, err)
					} else {
						log.Printf("[server] client %s closed", id)
					}
					return
				}
				line := reader.Text()
				log.Printf("[server] recv from %s: %q", id, line)
				h.submit(fmt.Sprintf("[broadcast] %s says: %s\n", id, line))
			}
		}(id, sc)
	}

	// 注：示例没有捕获信号关闭，为演示重点放在并发安全逻辑。
	cancel()
	wg.Wait()
}
